import chainlit as cl
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_ext.models.openai import OpenAIChatCompletionClient


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    if (username, password) == ("admin", "admin"):
        return cl.User(
            identifier="admin", metadata={"role": "admin", "provider": "credentials"}
        )
    elif (username, password) == ("puhaiyang", "123456"):
        return cl.User(
            identifier="puhaiyang", metadata={"role": "admin", "provider": "credentials"}
        )
    else:
        return None


@cl.on_chat_start
async def main():
    await cl.Message(content="您好，这里是牙体牙髓科，您牙齿哪里不适？").send()


# 此处定义的是FunctionCall
async def x_p_search(tooth_position: str) -> str:
    """Find information on the web"""
    app_user = cl.user_session.get("user")
    print(f"模拟查询{app_user.identifier}的{tooth_position}牙片数据")
    if tooth_position == "46":
        return "牙根尖处有阴影，疑似感染，需要进一步分析诊断"
    else:
        return f"{tooth_position}无影像"


async def run_team(query: str):
    model_client = OpenAIChatCompletionClient(model="deepseek-chat", base_url="https://api.deepseek.com",
                                              api_key="sk-6d775065fa8c405caebed674c125e67e",
                                              model_info={
                                                  "vision": False,
                                                  "function_calling": True,
                                                  "json_output": True,
                                                  "family": "unknown",
                                              })
    assistant_agent = AssistantAgent("assistant", model_client=model_client, tools=[x_p_search],
                                     system_message="你是一个牙体牙髓科的病情诊断机器人，负责对用户输入的症状描述分析原因，在分析病因前先询问出用户是具体哪一颗牙齿需要治疗。"
                                                    "在知道了具体的牙位号后，再调用x_p_search工具进行问题回答，传入给x_p_search工具的参数需要自动转为牙位号，如:28"
                                                    "如果用户的问题与病情咨询无关，回答：“您的描述与症状无关，暂不支持”")
    team = RoundRobinGroupChat(participants=[assistant_agent], max_turns=1)
    response_stream = team.run_stream(task=query)
    async for msg in response_stream:
        if hasattr(msg, "source") and (msg.type == "ToolCallExecutionEvent" or msg.type == "ToolCallRequestEvent"):
            # functionCall事件消息不显示给用户
            continue
        if hasattr(msg, "source") and msg.source != "user" and hasattr(msg, "content"):
            if msg.content.endswith("无影像"):
                res = await cl.AskActionMessage(
                    content=f"{msg.content}，是否需要帮您申请拍摄此牙的CT影像？",
                    actions=[
                        cl.Action(name="continue", payload={"value": "申请"}, label="✅ 申请牙片"),
                        cl.Action(name="cancel", payload={"value": "取消"}, label="❌ 取消"),
                    ],
                ).send()

                if res and res.get("payload").get("value") == "申请":
                    await cl.Message(
                        content="牙片申请已提交!待审核通过后前往第3影像室进行拍摄。",
                    ).send()
            else:
                msg = cl.Message(content=msg.content, author="Agent Team")
                await msg.send()


@cl.on_message
async def main(message: cl.Message):
    await run_team(message.content)
